package lsm

import (
	"bufio"
	"bytes"
	"io/fs"
	"kvdb/pkg/filter"
	"os"
	"path"
	"strings"
	"sync"
	"sync/atomic"
)

// lsm 主干架构：
// Conf 类聚合了有关于 lsm tree 运行所需的所有配置项，各项配置参数可以通过 option 模式进行灵活注入，对应的参数含义展示如下：
type Config struct {
	Dir      string // sst 文件存放的目录
	MaxLevel int    // lsm tree 总共多少层

	// sst 相关
	SSTSize          uint64 // 每个 sst table 大小，默认 4M
	SSTNumPerLevel   int    // 每层多少个 sstable，默认 10 个
	SSTDataBlockSize int    // sst table 中 block 大小 默认 16KB
	SSTFooterSize    int    // sst table 中 footer 部分大小. 固定为 32B

	Filter              filter.Filter                // 过滤器. 默认使用布隆过滤器
	MemTableConstructor memtable.MemTableConstructor // memtable 构造器，默认为跳表
}

// Tree 类对应的就是 lsm tree 的宏观结构：
type Tree struct {
	conf *Config

	// 读写数据时使用的锁
	dataLock sync.RWMutex

	// 每层 node 节点使用的读写锁
	levelLocks []sync.RWMutex

	// 读写 memtable
	memTable memtable.MemTable

	// 只读 memtable
	rOnlyMemTable []*memTableCompactItem

	// 预写日志写入口
	walWriter *wal.WALWriter

	// lsm树状数据结构
	nodes [][]*Node

	// memtable 达到阈值时，通过该 chan 传递信号，进行溢写工作
	memCompactC chan *memTableCompactItem

	// 某层 sst 文件大小达到阈值时，通过该 chan 传递信号，进行溢写工作
	levelCompactC chan int

	// lsm tree 停止时通过该 chan 传递信号
	stopc chan struct{}

	// memtable index，需要与 wal 文件一一对应
	memTableIndex int
	// 各层 sstable 文件 seq. sstable 文件命名为 level_seq.sst
	levelToSeq []atomic.Int32
}

// Node 类对应为 lsm tree 中的一个 sstable 文件，核心字段包括：
type Node struct {
	conf          *Config           // 配置文件
	file          string            // sstable 对应的文件名，不含目录路径
	level         int               // sstable 所在 level 层级
	seq           int32             // sstable 的 seq 序列号. 对应为文件名中的 level_seq.sst 中的 seq
	size          uint64            // sstable 的大小，单位 byte
	blockToFilter map[uint64][]byte // 各 block 对应的 filter bitmap
	index         []*Index          // 各 block 对应的索引
	startKey      []byte            // sstable 中最小的 key
	endKey        []byte            // sstable 中最大的 key
	sstReader     *SSTReader        // 读取 sst 文件的 reader 入口
}

// 每个 sstable 文件均以数据块 block 的形式进行数据分组. 后续无论是索引 index 还是过滤器 filter，都是与 block 对应存在的.
type Block struct {
	conf       *Config
	buffer     [30]byte      // 进行数据转移时使用的临时缓冲区
	record     *bytes.Buffer // 记录全量数据的缓冲区
	entriesCnt int           // kv 对数量
	prevKey    []byte        // 最晚一笔写入的数据的 key
}

// 对应于 lsm tree 中的一个 sstable. 这是写入流程的视角
type SSTWriter struct {
	conf          *Config           // 配置文件
	dest          *os.File          // sstable 对应的磁盘文件
	dataBuf       *bytes.Buffer     // 数据块缓冲区 key -> val
	filterBuf     *bytes.Buffer     // 过滤器块缓冲区 prev block offset -> filter bit map
	indexBuf      *bytes.Buffer     // 索引块缓冲区 index key -> prev block offset, prev block size
	blockToFilter map[uint64][]byte // prev block offset -> filter bit map
	index         []*Index          // index key -> prev block offset, prev block size

	dataBlock     *Block   // 数据块
	filterBlock   *Block   // 过滤器块
	indexBlock    *Block   // 索引块
	assistScratch [20]byte // 用于在写索引块时临时使用的辅助缓冲区

	prevKey         []byte // 前一笔数据的 key
	prevBlockOffset uint64 // 前一个数据块的起始偏移位置
	prevBlockSize   uint64 // 前一个数据块的大小
}

// 对应于 lsm tree 中的一个 sstable. 这是读取流程的视角
type SSTReader struct {
	conf         *Config       // 配置文件
	src          *os.File      // 对应的文件
	reader       *bufio.Reader // 读取文件的 reader
	filterOffset uint64        // 过滤器块起始位置在 sstable 的 offset
	filterSize   uint64        // 过滤器块的大小，单位 byte
	indexOffset  uint64        // 索引块起始位置在 sstable 的 offset
	indexSize    uint64        // 索引块的大小，单位 byte
}

// sstable 中用于快速检索 block 的索引
type Index struct {
	Key             []byte // 索引的 key. 保证其 >= 前一个 block 最大 key； < 后一个 block 的最小 key
	PrevBlockOffset uint64 // 索引前一个 block 起始位置在 sstable 中对应的 offset
	PrevBlockSize   uint64 // 索引前一个 block 的大小，单位 byte
}

// memtable 构造器
type MemTableConstructor func() MemTable

// 有序表 interface
// MemTable 是 lsm tree 中使用的内存有序表，此处声明为一个 interface，使用方可以根据需要实现具体的版本进行注入，默认会使用 golsm 项目下实现的跳表结构.
type MemTable interface {
	Put(key, value []byte)         // 写入数据
	Get(key []byte) ([]byte, bool) // 读取数据，第二个 bool flag 标识数据是否存在
	All() []*KV                    // 返回所有的 kv 对数据
	Size() int                     // 有序表内数据大小，单位 byte
	EntriesCnt() int               // kv 对数量
}

// 过滤器. 用于辅助 sstable 快速判定一个 key 是否存在于某个 block 中
type Filter interface {
	Add(key []byte)                // 添加 key 到过滤器
	Exist(bitmap, key []byte) bool // 是否存在 key
	Hash() []byte                  // 生成过滤器对应的 bitmap
	Reset()                        // 重置过滤器
	KeyLen() int                   // 存在多少个 key
}

/*
	启动流程
*/

// 配置文件构造器.
func NewConfig(dir string, opts ...ConfigOption) (*Config, error) {
	c := Config{
		Dir:           dir, // sstable 文件所在的目录路径
		SSTFooterSize: 32,  // 对应 4 个 uint64，共 32 byte
	}
	// 加载配置项
	for _, opt := range opts {
		opt(&c)
	}
	// 兜底修复
	repaire(&c)
	return &c, c.check() // 校验一下配置是否合法，主要是 check 存放 sst 文件和 wal 文件的目录，如果有缺失则进行目录创建
}

// 校验一下配置是否合法，主要是 check 存放 sst 文件和 wal 文件的目录，如果有缺失则进行目录创建
func (c *Config) check() error {
	// sstable 文件目录确保存在
	if _, err := os.ReadDir(c.Dir); err != nil {
		_, ok := err.(*fs.PathError)
		if !ok || !strings.HasSuffix(err.Error(), "no such file or directory") {
			return err
		}
		if err = os.Mkdir(c.Dir, os.ModePerm); err != nil {
			return err
		}
	}
	// wal 文件目录确保存在
	walDir := path.Join(c.Dir, "walfile")
	if _, err := os.ReadDir(walDir); err != nil {
		_, ok := err.(*fs.PathError)
		if !ok || !strings.HasSuffix(err.Error(), "no such file or directory") {
			return err
		}
		if err = os.Mkdir(walDir, os.ModePerm); err != nil {
			return err
		}
	}

	return nil
}

// 配置项
type ConfigOption func(*Config)

// lsm tree 最大层数. 默认为 7 层.
func WithMaxLevel(maxLevel int) ConfigOption {
	return func(c *Config) {
		c.MaxLevel = maxLevel
	}
}

// level0层每个 sstable 文件的大小，单位 byte. 默认为 1 MB.
// 且每加深一层，sstable 文件大小限制阈值放大 10 倍.
func WithSSTSize(sstSize uint64) ConfigOption {
	return func(c *Config) {
		c.SSTSize = sstSize
	}
}

// sstable 中每个 block 块的大小限制. 默认为 16KB.
func WithSSTDataBlockSize(sstDataBlockSize int) ConfigOption {
	return func(c *Config) {
		c.SSTDataBlockSize = sstDataBlockSize
	}
}

// 每个 level 层预期最多存放的 sstable 文件个数. 默认为 10 个.
func WithSSTNumPerLevel(sstNumPerLevel int) ConfigOption {
	return func(c *Config) {
		c.SSTNumPerLevel = sstNumPerLevel
	}
}

// 注入过滤器的具体实现. 默认使用本项目下实现的布隆过滤器 bloom filter.
func WithFilter(filter filter.Filter) ConfigOption {
	return func(c *Config) {
		c.Filter = filter
	}
}

// 注入有序表构造器. 默认使用本项目下实现的跳表 skiplist.
func WithMemtableConstructor(memtableConstructor memtable.MemTableConstructor) ConfigOption {
	return func(c *Config) {
		c.MemTableConstructor = memtableConstructor
	}
}

func repaire(c *Config) {
	// lsm tree 默认为 7 层.
	if c.MaxLevel <= 1 {
		c.MaxLevel = 7
	}

	// level0 层每个 sstable 文件默认大小限制为 1MB.
	// 且每加深一层，sstable 文件大小限制阈值放大 10 倍.
	if c.SSTSize <= 0 {
		c.SSTSize = 1024 * 1024
	}

	// sstable 中每个 block 块的大小限制. 默认为 16KB.
	if c.SSTDataBlockSize <= 0 {
		c.SSTDataBlockSize = 16 * 1024 // 16KB
	}

	// 每个 level 层预期最多存放的 sstable 文件个数. 默认为 10 个.
	if c.SSTNumPerLevel <= 0 {
		c.SSTNumPerLevel = 10
	}

	// 注入过滤器的具体实现. 默认使用本项目下实现的布隆过滤器 bloom filter.
	if c.Filter == nil {
		c.Filter, _ = filter.NewBloomFilter(1024)
	}

	// 注入有序表构造器. 默认使用本项目下实现的跳表 skiplist.
	if c.MemTableConstructor == nil {
		c.MemTableConstructor = memtable.NewSkiplist
	}
}

// 构建出一棵 lsm tree
func NewTree(conf *Config) (*Tree, error) {
	// 1 构造 lsm tree 实例
	t := Tree{
		conf:          conf,
		memCompactC:   make(chan *memTableCompactItem),
		levelCompactC: make(chan int),
		stopc:         make(chan struct{}),
		levelToSeq:    make([]atomic.Int32, conf.MaxLevel),
		nodes:         make([][]*Node, conf.MaxLevel),
		levelLocks:    make([]sync.RWMutex, conf.MaxLevel),
	}

	// 2 读取 sst 文件，还原出整棵树
	if err := t.constructTree(); err != nil {
		return nil, err
	}

	// 3 运行 lsm tree 压缩调整协程
	go t.compact()

	// 4 读取 wal 还原出 memtable
	if err := t.constructMemtable(); err != nil {
		return nil, err
	}

	// 5 返回 lsm tree 实例
	return &t, nil
}

// constructTree 方法用于读取已有的 sst 文件，在内存中还原出一系列 lsm node 节点// 读取 sst 文件，还原出整棵树
func (t *Tree) constructTree() error {
	// 读取 sst 文件目录下的 sst 文件列表
	sstEntries, err := t.getSortedSSTEntries()
	if err != nil {
		return err
	}

	// 遍历每个 sst 文件，将其加载为 node 添加 lsm tree 的 nodes 内存切片中
	for _, sstEntry := range sstEntries {
		if err = t.loadNode(sstEntry); err != nil {
			return err
		}
	}

	return nil
}

// constructMemTable 方法用于读取已有 wal 文件，在内存中还原出一系列只读 memtable 和读写 memtable// 读取 wal 还原出 memtable
func (t *Tree) constructMemtable() error {
	// 1 读 wal 目录，获取所有的 wal 文件
	wals, _ := os.ReadDir(path.Join(t.conf.Dir, "walfile"))

	// 2 倘若 wal 目录不存在或者 wal 文件不存在，则构造一个新的 memtable
	if len(wals) == 0 {
		t.newMemTable()
		return nil
	}

	// 3 依次还原 memtable. 最晚一个 memtable 作为读写 memtable
	// 前置 memtable 作为只读 memtable，分别添加到内存 slice 和 channel 中.
	return t.restoreMemTable(wals)
}

// 运行 compact 协程.
func (t *Tree) compact() {
	for {
		select {
		// 接收到 lsm tree 终止信号，退出协程.
		case <-t.stopc:
			// log
			return
			// 接收到 read-only memtable，需要将其溢写到磁盘成为 level0 层 sstable 文件.
		case memCompactItem := <-t.memCompactC:
			t.compactMemTable(memCompactItem)
			// 接收到 level 层 compact 指令，需要执行 level~level+1 之间的 level sorted merge 流程.
		case level := <-t.levelCompactC:
			t.compactLevel(level)
		}
	}
}

/*
	写流程
*/
//将 kv 对数据写入预写日志中，防止 memtable 丢失数据
//将 kv 对数据写入读写 memtable
// 写入一组 kv 对到 lsm tree. 会直接写入到读写 memtable 中.
func (t *Tree) Put(key, value []byte) error {
	// 1 加写锁
	t.dataLock.Lock()
	defer t.dataLock.Unlock()

	// 2 数据预写入预写日志中，防止因宕机引起 memtable 数据丢失.
	if err := t.walWriter.Write(key, value); err != nil {
		return err
	}

	// 3 数据写入读写跳表
	t.memTable.Put(key, value)

	// 4 倘若读写跳表的大小未达到 level0 层 sstable 的大小阈值，则直接返回.
	// 考虑到溢写成 sstable 后，需要有一些辅助的元数据，预估容量放大为 5/4 倍
	if uint64(t.memTable.Size()*5/4) <= t.conf.SSTSize {
		return nil
	}

	// 5 倘若读写跳表数据量达到上限，则需要切换跳表
	t.refreshMemTableLocked()
	return nil
}

// 值得一提的是，在写流程结束前，需要校验一下读写 memtable 中的数据量大小是否达到溢写落盘的阈值，是的话，需要通过 refreshMemTableLocked 方法完成 memtable 的模式切换：将 memtable 由读写模式切换为只读模式，然后通过 memCompactC 通道将其传递给 compact 协程进行溢写处理构建一个新的 memtable 用于处理写请求，并构造与之相应的 wal 文件// 切换读写跳表为只读跳表，并构建新的读写跳表
func (t *Tree) refreshMemTableLocked() {
	// 辞旧
	// 将读写跳表切换为只读跳表，追加到 slice 中，并通过 chan 发送给 compact 协程，由其负责进行溢写成为 level0 层 sst 文件的操作.
	oldItem := memTableCompactItem{
		walFile:  t.walFile(),
		memTable: t.memTable,
	}
	t.rOnlyMemTable = append(t.rOnlyMemTable, &oldItem)
	t.walWriter.Close()
	go func() {
		t.memCompactC <- &oldItem
	}()

	// 迎新
	// 构造一个新的读写 memtable，并构造与之相应的 wal 文件.
	t.memTableIndex++
	t.newMemTable()
}

// 读流程
// 根据 key 读取数据.
func (t *Tree) Get(key []byte) ([]byte, bool, error) {
	t.dataLock.RLock()
	// 1 首先读 active memtable.
	value, ok := t.memTable.Get(key)
	if ok {
		t.dataLock.RUnlock()
		return value, true, nil
	}

	// 2 读 readOnly memtable
	for i := len(t.rOnlyMemTable) - 1; i >= 0; i-- {
		value, ok = t.rOnlyMemTable[i].memTable.Get(key)
		if ok {
			t.dataLock.RUnlock()
			return value, true, nil
		}
	}
	t.dataLock.RUnlock()

	// 3 读 sstable level0 层.
	var err error
	t.levelLocks[0].RLock()
	for i := len(t.nodes[0]) - 1; i >= 0; i-- {
		if value, ok, err = t.nodes[0][i].Get(key); err != nil {
			t.levelLocks[0].RUnlock()
			return nil, false, err
		}
		if ok {
			t.levelLocks[0].RUnlock()
			return value, true, nil
		}
	}
	t.levelLocks[0].RUnlock()

	// 4 依次读 sstable level 1 ~ i 层.
	for level := 1; level < len(t.nodes); level++ {
		t.levelLocks[level].RLock()
		node, ok := t.levelBinarySearch(level, key, 0, len(t.nodes[level])-1)
		if !ok {
			t.levelLocks[level].RUnlock()
			continue
		}
		if value, ok, err = node.Get(key); err != nil {
			t.levelLocks[level].RUnlock()
			return nil, false, err
		}
		if ok {
			t.levelLocks[level].RUnlock()
			return value, true, nil
		}
		t.levelLocks[level].RUnlock()
	}

	// 5 至此都没有读到数据，则返回 key 不存在.
	return nil, false, nil
}
